"""
AWS instance management for Dumbo benchmark
"""

import boto3
import time
from typing import List, Dict, Any
from benchmark.utils import BenchError, Print


class AWSInstanceManager:
    """AWS instance manager for Dumbo benchmark"""
    
    def __init__(self):
        self.ec2_client = boto3.client('ec2')
        self.instances = []
    
    def create_instances(self, nodes: int, instance_type: str = 't3.medium', region: str = 'us-east-1') -> List[Dict[str, Any]]:
        """Create AWS instances for Dumbo benchmark"""
        Print.info(f"Creating {nodes} AWS instances of type {instance_type} in {region}")
        
        try:
            # Create instances
            response = self.ec2_client.run_instances(
                ImageId='ami-0c02fb55956c7d316',  # Amazon Linux 2 AMI
                MinCount=nodes,
                MaxCount=nodes,
                InstanceType=instance_type,
                KeyName='dumbo-aws-key',  # You need to create this key pair
                SecurityGroupIds=['sg-12345678'],  # You need to create this security group
                TagSpecifications=[
                    {
                        'ResourceType': 'instance',
                        'Tags': [
                            {'Key': 'Name', 'Value': 'dumbo-benchmark'},
                            {'Key': 'Project', 'Value': 'dumbo'},
                            {'Key': 'Type', 'Value': 'benchmark'}
                        ]
                    }
                ]
            )
            
            # Store instance information
            for instance in response['Instances']:
                self.instances.append({
                    'id': instance['InstanceId'],
                    'type': instance_type,
                    'region': region,
                    'state': 'pending'
                })
            
            Print.info(f"Created {len(self.instances)} instances")
            
            # Wait for instances to be running
            self._wait_for_instances()
            
            return self.instances
            
        except Exception as e:
            raise BenchError(f"Failed to create AWS instances: {e}")
    
    def _wait_for_instances(self):
        """Wait for all instances to be running"""
        Print.info("Waiting for instances to be running...")
        
        instance_ids = [inst['id'] for inst in self.instances]
        
        waiter = self.ec2_client.get_waiter('instance_running')
        try:
            waiter.wait(InstanceIds=instance_ids)
            
            # Get updated instance information
            response = self.ec2_client.describe_instances(InstanceIds=instance_ids)
            for reservation in response['Reservations']:
                for instance in reservation['Instances']:
                    for i, inst in enumerate(self.instances):
                        if inst['id'] == instance['InstanceId']:
                            self.instances[i]['public_ip'] = instance.get('PublicIpAddress', '')
                            self.instances[i]['private_ip'] = instance.get('PrivateIpAddress', '')
                            self.instances[i]['state'] = instance['State']['Name']
                            break
            
            Print.success("All instances are running")
            
        except Exception as e:
            raise BenchError(f"Failed to wait for instances: {e}")
    
    def get_instance_info(self) -> List[Dict[str, Any]]:
        """Get information about all instances"""
        return self.instances
    
    def terminate_instances(self):
        """Terminate all instances"""
        if not self.instances:
            Print.info("No instances to terminate")
            return
        
        Print.info("Terminating AWS instances...")
        
        instance_ids = [inst['id'] for inst in self.instances]
        
        try:
            self.ec2_client.terminate_instances(InstanceIds=instance_ids)
            Print.success("All instances terminated")
            
            # Clear instances list
            self.instances = []
            
        except Exception as e:
            raise BenchError(f"Failed to terminate instances: {e}")
    
    def create_security_group(self, group_name: str = 'dumbo-benchmark-sg') -> str:
        """Create security group for Dumbo benchmark"""
        Print.info(f"Creating security group: {group_name}")
        
        try:
            # Create security group
            response = self.ec2_client.create_security_group(
                GroupName=group_name,
                Description='Security group for Dumbo benchmark'
            )
            
            group_id = response['GroupId']
            
            # Add inbound rules
            self.ec2_client.authorize_security_group_ingress(
                GroupId=group_id,
                IpPermissions=[
                    {
                        'IpProtocol': 'tcp',
                        'FromPort': 22,
                        'ToPort': 22,
                        'IpRanges': [{'CidrIp': '0.0.0.0/0'}]
                    },
                    {
                        'IpProtocol': 'tcp',
                        'FromPort': 9000,
                        'ToPort': 9100,
                        'IpRanges': [{'CidrIp': '0.0.0.0/0'}]
                    }
                ]
            )
            
            Print.success(f"Security group created: {group_id}")
            return group_id
            
        except Exception as e:
            raise BenchError(f"Failed to create security group: {e}")
    
    def create_key_pair(self, key_name: str = 'dumbo-aws-key') -> str:
        """Create key pair for AWS instances"""
        Print.info(f"Creating key pair: {key_name}")
        
        try:
            response = self.ec2_client.create_key_pair(KeyName=key_name)
            
            # Save private key to file
            key_file = f"{key_name}.pem"
            with open(key_file, 'w') as f:
                f.write(response['KeyMaterial'])
            
            # Set proper permissions
            import os
            os.chmod(key_file, 0o400)
            
            Print.success(f"Key pair created: {key_file}")
            return key_file
            
        except Exception as e:
            raise BenchError(f"Failed to create key pair: {e}")

